# encoding: utf-8
# 数据集：http://s3.bmio.net/kashgari/china-people-daily-ner-corpus.tar.gz


from lsptrain.nlp_ner.train_scripts import train


if __name__ == '__main__':
    pretrained_model = r"D:/codes/nlp_about/pretrained_model/hfl_chinese-roberta-wwm-ext"
    # train_path = r"C:\Users\Administrator\Downloads\china-people-daily-ner-corpus\china-people-daily-ner-corpus\renminribao.txt"
    train_path = r"C:\Users\Administrator\Downloads\china-people-daily-ner-corpus\china-people-daily-ner-corpus\example.train"
    labels = ["LOC", "ORG", "PER"]

    train(train_path, labels,
          pretrained_model=pretrained_model,
          batch_size=16,
          save_best=False,
          max_length=64,
          num_epochs=5

          )